#include <cstdlib>
#include <cstring>
#include <iostream>
#include <fstream>
#include <glog/logging.h>
//jsoncpp 相关的头文件
#include "reader.h"
#include "writer.h"
#include "value.h"
#include "SampleDetector.hpp"

SampleDetector::SampleDetector()
{
    
}

SampleDetector::~SampleDetector()
{
    UnInit();
}

STATUS SampleDetector::Init(const std::string &model_path, float thresh)
{        
    Core core;
    model = core.read_model(model_path);
    compiled_model = core.compile_model(model, device_name);
    infer_request = compiled_model.create_infer_request();
    input_tensor = infer_request.get_input_tensor(0);
    
    string classesFile = "/usr/local/ev_sdk/model/coco.names";
    ifstream ifs(classesFile.c_str());
	string line;
	while (getline(ifs, line)) this->classes.push_back(line);
    num_class = classes.size();

    mThresh = thresh;
    nms_area_threshold = 0.5;
    SDKLOG(INFO) << "Init Done. model name is " << model_path << ",thresh is " << mThresh;
    return 0;
}

STATUS SampleDetector::UnInit()
{
    if(model != nullptr)
    {
        /*delete model;
        model = nullptr;*/
        model.reset();
    }
    return STATUS_SUCCESS;
}

/**
 * @brief demo自定义辅助函数
 */
void SampleDetector::nms(vector<Object>& inputBoxes)
{
    std::sort(inputBoxes.begin(), inputBoxes.end(), [](Object a, Object b) { return a.prob > b.prob; });

    std::vector<bool> isSuppressed(inputBoxes.size(), false);
    for (int i = 0; i < int(inputBoxes.size()); ++i)
    {
        if (isSuppressed[i])
        {
            continue;
        }
        for (int j = i + 1; j < int(inputBoxes.size()); ++j)
        {
            if (isSuppressed[j])
            {
                continue;
            }

            float inter = (inputBoxes[i].rect&inputBoxes[j].rect).area();
            float ovr = inter / (inputBoxes[i].rect.area() + inputBoxes[j].rect.area() - inter);

            if (ovr >= this->nms_area_threshold)
            {
                isSuppressed[j] = true;
            }
        }
    }

    // return post_nms;
    int idx_t = 0;
    inputBoxes.erase(remove_if(inputBoxes.begin(), inputBoxes.end(), [&idx_t, &isSuppressed](const Object &f) { return isSuppressed[idx_t++]; }), inputBoxes.end());
}

STATUS SampleDetector::ProcessImage(cv::Mat &inFrame, std::vector<Object> &result, float thresh)
{
    mThresh = thresh;
    if (inFrame.empty())
    {
        SDKLOG(ERROR) << "Invalid input!";
        return -1;
    }
    result.clear();
    
    float r = std::min(mInpHeight / static_cast<float>(inFrame.rows), mInpWidth / static_cast<float>(inFrame.cols));
    cv::Size new_size = cv::Size{inFrame.cols * r, inFrame.rows * r};    
    Mat resize_img;
    resize(inFrame, resize_img, new_size);
    cvtColor(resize_img, resize_img, COLOR_BGR2RGB);
    Mat dstimg = cv::Mat(cv::Size(mInpWidth, mInpHeight), CV_8UC3, cv::Scalar(114, 114, 114));    
    resize_img.copyTo(dstimg(cv::Rect{0, 0, resize_img.cols, resize_img.rows}));
    
    size_t img_area = mInpHeight*mInpWidth;
    auto blob_data = input_tensor.data<float>();
    //nchw
    for(size_t row =0;row<mInpHeight;row++){
        for(size_t col=0;col<mInpWidth;col++){
            for(size_t ch =0;ch<3;ch++){
                blob_data[img_area*ch + row*mInpWidth + col] = float(dstimg.at<Vec3b>(row,col)[ch])/255.0f;
            }
        }
    }
    //执行预测
    infer_request.infer();
    auto output_tensor = infer_request.get_output_tensor(0);
    const float *outs = output_tensor.data<const float>();
    
    decode_outputs(outs, result, 1/r, inFrame.cols, inFrame.rows);
    //后处理获得最终检测结果
    this->nms(result); 
    return 0;
}

bool SampleDetector::decode_outputs(const float *outs, vector<Object>& detected_objects, const float scale, const int img_w, const int img_h)
{
    const int box_num = 80*80*3+40*40*3+20*20*3;
    for(int n=0;n<box_num;n++)
    {
        const float* pbuf = outs + n*(5+num_class);
        float box_prob = pbuf[4];
        if(box_prob > mThresh)
        {
            float max_prob = 0;
            int idx=0;
            for(int i=0;i<num_class;i++)
            {
                if(pbuf[i+5] > max_prob)
                {
                    max_prob = pbuf[i+5];
                    idx = i;
                }
            }
            float score = box_prob * max_prob;
            if(score > mThresh)
            {
                float x = pbuf[0];
                float y = pbuf[1];
                float w = pbuf[2];
                float h = pbuf[3];
                
                x *= scale;
                y *= scale;
                w *= scale;
                h *= scale;

                float r_x = x - w*0.5;
                float r_y = y - h*0.5;
                Rect box = Rect(round(r_x),round(r_y),round(w),round(h));
                box = box & cv::Rect(0, 0, img_w-1, img_h-1);
                if( box.area() > 0)
                {
                    detected_objects.push_back({score, this->classes[idx], box});
                } 
            } 
        }
    }
    if(detected_objects.size() == 0) return false;
    else return true;
}

